# -*- coding: utf-8 -*-

"""Datastore Profiler"""

# See README on how to use this module.

import logging
import os
import sys
import wsgiref.handlers
from cStringIO import StringIO
from functools import partial

from google.appengine.api import apiproxy_stub_map
from google.appengine.api import memcache
from google.appengine.ext import webapp
from google.appengine.ext.webapp.template import Template, Context

from django.utils.html import escape


IGNORE_CALLS=['Next']      # List of calls to ignore when logging.
OUTPUT_CHANNEL='memcache'  # allowed: 'inline', 'memcache', 'log'.


# No changes needed below this line...

DEBUG_QUERIES = []


def LoggingMakeSyncCall(real_func, service, call, request, response):
  """Wrapper for apiproxy_stub_map.MakeSyncCall."""
  global IGNORE_CALLS, DEBUG_QUERIES
  def get_summary(obj):
    from google.appengine.datastore import datastore_pb
    if isinstance(obj, datastore_pb.Query):
      ret = 'Kind: %s' % obj.kind()
      if obj.filter_size() > 0:
        ret += '; Filter: %s' % (', '.join(str(x) for x in obj.filter_list()))
      return ret
    elif isinstance(obj, datastore_pb.QueryResult):
      if obj.has_cursor():
        return str(obj.cursor())
      else:
        return 'QueryResult: %d results' % obj.result_size()
    elif isinstance(obj, datastore_pb.GetRequest):
      return 'Keys: %s' % (', '.join(str(x) for x in obj.key_list()))
    elif isinstance(obj, datastore_pb.GetResponse):
      return '%d entities' % obj.entity_size()
    else:
      lines = str(obj).splitlines()
      if len(lines) > 0:
        return lines[0]
      else:
        return '???'

  # Check some variables. Sometimes they become invalid when SDK reloads.
  if IGNORE_CALLS is None:
    IGNORE_CALLS = ['Next']
  if DEBUG_QUERIES is None:
    DEBUG_QUERIES = []
  import time
  start = time.time()
  ret = real_func(service, call, request, response)
  exectime = time.time()-start
  if service == 'datastore_v3':
    if call not in IGNORE_CALLS:
      str_request = str(request)[:1000000]
      str_response = str(response)[:1000000]
      DEBUG_QUERIES.append((call, exectime,
                            (get_summary(request), str_request),
                            (get_summary(response), str_response),
                             ))
  return ret

# Let's patch MakeSyncCall
apiproxy_stub_map.MakeSyncCall = partial(LoggingMakeSyncCall,
                                         apiproxy_stub_map.MakeSyncCall)


def _debug_main(real_main):
  """Main program for debugging queries."""
  global DEBUG_QUERIES, OUTPUT_CHANNEL
  if DEBUG_QUERIES is None:
    DEBUG_QUERIES = []
  # Replace stdout to catch response code
  old_stdout = sys.stdout
  tmp_out = StringIO()
  sys.stdout = tmp_out
  # Run main function
  real_main()
  # Analyze results
  content = tmp_out.getvalue()
  if content.startswith('Status: 200 OK'): # keep data when redirecting
    if OUTPUT_CHANNEL == 'inline':
      debug_out = _data2inline(DEBUG_QUERIES)
      content = content.replace('</body>', '%s\n</body>' % ''.join(debug_out))
    elif OUTPUT_CHANNEL == 'log':
      debug_out = _data2log(DEBUG_QUERIES)
      logging.debug('DATASTORE:\n%s' % '\n'.join(debug_out))
    elif OUTPUT_CHANNEL == 'memcache':
      keys = memcache.get('datastoredbg_keys')
      if keys is None:
        keys = []
      nums = [x[1] for x in keys]
      nums.sort()
      if len(nums) == 0:
        num = 0
      else:
        num = nums[-1]
      num += 1
      memcache.set('dbg_%d' % num, DEBUG_QUERIES)
      keys.append(('%s %s' % (os.environ['REQUEST_METHOD'],
                              os.environ['PATH_INFO']),
                   num))
      memcache.set('datastoredbg_keys', keys)
    DEBUG_QUERIES = []
  sys.stdout = old_stdout
  sys.stdout.write(content)
  sys.stdout.flush()


def debug_main(real_main):
  """Returns a wrapper for our main function."""
  return partial(_debug_main, real_main)


def _data2inline(data):
  """Converts the collected data into a HTML snippet."""
  total_time = 0
  debug_out = ['\n<hr>']
  debug_out.append('<h1>Queries</h1>')
  for action, exectime, req, resp in data:
    debug_out.append('<b>Action:')
    debug_out.append('%s (%.4f secs)</b><br/>' % (action, exectime))
    debug_out.append('<i>Request:</i> %s<br/>' % escape(req[0]))
    # Uncomment the following two lines if you want to see the response.
    # Note: The response could be very huge...
    #debug_out.append('<br/>')
    #debug_out.append('<i>Response:</i> %s' % escape(resp[0]))
    debug_out.append('<br/><br/>')
    total_time += exectime
  debug_out.insert(2,
                   '<b>Datastore accessed %d times.</b><br/><br/>'
                   % len(data))
  debug_out.insert(2, '<b>Total time: %.4f secs.</b><br/>' % total_time)
  return debug_out


def _data2log(data):
  """Converts collected data to plain text suitable for logging."""
  debug_out = []
  for action, exectime, req, resp in data:
    debug_out.append('Action: %s (%.4f secs)' % (action, exectime))
    debug_out.append(req)
    debug_out.append(resp)
  return debug_out


# Request handler for memcache viewer


class DatastoreDebugHandler(webapp.RequestHandler):
  """Displays collected data from memcache."""

  def get(self):
    if self.request.path.endswith('fetchData'):
      self._fetch_data()
    elif self.request.path.endswith('flush'):
      keys = memcache.get('datastoredbg_keys')
      if keys:
        nums = [x[1] for x in keys]
        delkeys = ['dbg_%d' % x[1] for x in keys]
        memcache.delete_multi(delkeys)
      memcache.set('datastoredbg_keys', [])
      self.redirect('/'.join(self.request.path.rsplit('/')[:-1])+'/')
    else:
      tpl = Template(TPL_MEMCACHE)
      ctx = Context({'memcachekeys': memcache.get('datastoredbg_keys')})
      self.response.out.write(tpl.render(ctx))

  def _fetch_data(self):
    num = self.request.get('num')
    data = memcache.get('dbg_%s' % num)
    if data is None:
      self.response.out.write('<span style="color:red;">Key not found.</span>')
    else:
      count = 0
      total_time = 0
      for action, exectime, req, resp in data:
        count += 1
        total_time += exectime
        self.response.out.write('<b>Action: %s (%.4f secs.)</b><br/>'
                                % (action, exectime))
        self.response.out.write(('<div><a class="more" '
                                 'href="javascript:tgl(\'details-req-%s-%s\')";'
                                 '><b>Request:</b> %s</a></div>'
                                 % (num, count, escape(req[0]))))
        details = escape(req[1]).replace('\n', '<br/>').replace(' ', '&nbsp;')
        self.response.out.write(('<div id="details-req-%s-%s" class="details" '
                                 'style="display:none;">%s</div>'
                                 % (num, count, details)))
        self.response.out.write(('<div><a class="more" '
                                 'href="javascript:tgl(\'details-res-%s-%s\')";'
                                 '><b>Response:</b> %s</a></div>'
                                 % (num, count, escape(resp[0]))))
        details = escape(resp[1]).replace('\n', '<br/>').replace(' ', '&nbsp;')
        self.response.out.write(('<div id="details-res-%s-%s" class="details" '
                                 'style="display:none;">%s</div>'
                               % (num, count, details)))
        self.response.out.write('<hr/>')
      self.response.out.write('<b>Total time: %.4f secs.</b><br/>'
                              % total_time)
      self.response.out.write('<b>Datastore accessed %d times.</b>' % count)


def main():
  application = webapp.WSGIApplication(
    [('.*', DatastoreDebugHandler)])
  wsgiref.handlers.CGIHandler().run(application)


TPL_MEMCACHE ='''<html>
<head>
  <script src="http://www.google.com/jsapi"></script>
  <script>
    google.load("jquery", "1.2.6");
  </script>
  <script>
    function fetchData(num) {
      $.get("fetchData?num="+num, function(data, textStatus) {
        $("#data"+num).html(data);
        $("#action"+num).attr("href", "javascript:tgl(\'data"+num+"\');");
      });
    }
    function tgl(id) {
      $("#"+id).toggle();
    }
    function expandAll() {
      $(".linkaction").each(function(idx) {
        obj = $(this);
        action = obj.attr("href");
        match = action.match(/fetchData\((\d+)\)/);
        if ( match ) { fetchData(match[1]); }
      });
    }
  </script>
  <style>
    body { font-family: Arial, sans; font-size: 10pt; }
    .entry {
      margin-bottom: .5em;
      border: 1px solid #aaaaaa;
      padding: 3px;
      background: #dddddd;
    }
    .more {
      text-decoration: none;
      color: #336699;
      padding-left: .5em;
    }
    .details {
      margin-left: 1em;
      margin-right: 1em;
      background-color: white;
      border: 1px solid #aaaaaa;
      font-family: "Courier New", monospace;
      padding: 3px;
    }
  </style>
</head>
<body>
<h1>Datastore Access</h1>
<div style="margin-bottom: 1.5em;">
  <a href="javascript:expandAll();">Expand all</a>
  |
  <a href="flush">Flush cached data</a>
  |
  <a href="">Reload</a>
</div>
{% for mckey in memcachekeys %}
<div id="container{{mckey.1}}" class="entry">
  <a id="action{{mckey.1}}" class="linkaction"
     href="javascript:fetchData({{mckey.1}});">{{mckey.0}}</a>
  <div id="data{{mckey.1}}"></div>
</div>
{% endfor %}
</body>
</html>'''


if __name__ == '__main__':
  main()
